[rollout, diffusion] fix: support prompt embedding cache for Qwen-Image by Sky-Trigger · Pull Request #214 · verl-project/verl-omni

Sky-Trigger · 2026-06-29T06:21:49Z

Keep token IDs and attention masks hashable until encode_prompt(), convert them to tensors inside token-ID encoders, and expose vLLM-Omni prompt embedding cache controls through the diffusion rollout configuration.

What does this PR do?

This PR adapts verl-omni's tokenized diffusion prompt path to the prompt embedding cache introduced by vLLM-Omni PR #2962.

Previously, rollout adapters converted token ID lists and masks to torch.Tensor before calling encode_prompt(). Since the vLLM-Omni cache cannot build keys from tensor arguments, these calls bypassed the cache.

This change keeps prompt IDs and masks list-based across the cache boundary and converts them to tensors only inside the underlying token-ID encoder. It covers:

Qwen-Image FlowGRPO and MixGRPO
Qwen-Image DiffusionNFT
Qwen-Image online DPO
Qwen-Image stepwise execution
Wan2.2 DanceGRPO

It also adds rollout configuration for enabling and sizing the cache without requiring environment variables.

Checklist Before Starting

Searched for similar PRs: open PR search
Formatted the PR title as:
[rollout, diffusion] fix: support prompt embedding cache for tokenized prompts

Test

Static checks completed:

ruff check
ruff format --check
Python compileall
git diff --check

Manual validation:

examples/diffusionnft_trainer/run_qwen_image_ocr_lora.sh completed successfully with the updated DiffusionNFT adapter.

Configuration unit tests were added to tests/workers/config/test_diffusion_config_on_cpu.py.

Full pytest and all adapter-specific GPU/NPU experiments were not run in the current environment. The remaining independent rollout paths are FlowGRPO, online DPO, and Wan2.2 DanceGRPO.

API and Usage Example

The cache is disabled by default and uses the upstream default capacity of 32.

It can be enabled through Hydra configuration:

python3 -m verl_omni.trainer.main_diffusion \
    actor_rollout_ref.rollout.enable_prompt_embed_cache=True \
    actor_rollout_ref.rollout.prompt_embed_cache_size=64 \
    ...

Equivalent YAML:

actor_rollout_ref:
  rollout:
    enable_prompt_embed_cache: true
    prompt_embed_cache_size: 64

No example shell scripts are modified.

Design & Code Changes

Prompt cache boundary

Before:

adapter
  -> list to torch.Tensor
  -> cached encode_prompt()
  -> tensor argument is not hashable
  -> cache bypass

After:

adapter
  -> list-based prompt IDs and masks
  -> cached encode_prompt()
  -> cache lookup
  -> on miss, token-ID encoder converts lists to tensors
  -> text encoder

Adapter changes

Added list-compatible batch-size calculation.
Moved list-to-tensor conversion into Qwen-Image and Wan token-ID encoders.
Kept positive and negative prompt IDs hashable until encode_prompt().
Normalized rollout prompt_mask to a list before constructing the diffusion request.
Preserved the existing DPO and stepwise class structures and inheritance relationships.
Left SD3 and BAGEL unchanged because they do not use this token-ID cache path.

Configuration changes

Added to DiffusionRolloutConfig:

enable_prompt_embed_cache: bool = False
prompt_embed_cache_size: int = 32

These values are passed directly to AsyncOmni. Cache size must be a positive integer.

Checklist Before Submitting

Read the [Contribute Guide](https://github.com/verl-project/verl-omni/blob/main/CONTRIBUTING.md).
Run full pre-commit checks:
pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always
Add / Update documentation.
- No standalone documentation was added; the new configuration fields are documented in the rollout YAML.
Added configuration unit tests to the existing CPU config test suite.
- Adapter behavior was validated through DiffusionNFT end-to-end execution.
- Additional GPU/NPU coverage is not added to CI because it requires model weights and accelerator resources.

AI Assistance

AI assistance was used to inspect the adapter call paths, implement the changes, and prepare this PR description. All changed lines and validation results were reviewed by the human submitter.

gemini-code-assist

Code Review

This pull request introduces a prompt embedding cache mechanism for diffusion workers by adding enable_prompt_embed_cache and prompt_embed_cache_size configurations across trainer and rollout configs. It also refactors several rollout adapters (such as Qwen and Wan22) to support list-based inputs for prompt IDs and masks, standardizing batch size calculation via a new helper function get_prompt_batch_size. There are no review comments provided, so I have no feedback to provide.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

zhtmike · 2026-06-29T06:35:42Z

@knlnguyen1802 , prompt cache related, PTAL

knlnguyen1802 · 2026-07-03T08:59:20Z

@Sky-Trigger Could you help to run a verification and post a curve

Sky-Trigger · 2026-07-06T06:57:31Z

I have verified this change with DiffusionNFT, FlowGRPO, and related models. This modification only delays the point at which prompt_ids and other inputs are converted into tensors, so that the code path is compatible with vLLM-Omni PR #2962. It does not introduce any numerical changes, precision-related behavior, or new model structure.

knlnguyen1802 · 2026-07-06T09:13:30Z

@Sky-Trigger Thanks for verify it, could you run several step to compare performance gain on this branch ?

Sky-Trigger · 2026-07-07T01:23:34Z

I conducted tests using run_qwen_image_ocr_lora_npu.sh in FlowGRPO, comparing runs with and without prompt cache enabled. Each configuration was run for 30 steps. With prompt cache enabled, the average step time decreased from 432.28s to 423.76s, achieving an average reduction of 1.97%.

knlnguyen1802 · 2026-07-07T08:59:56Z

-        prompt_ids: torch.Tensor,
-        prompt_mask: torch.Tensor | None,
-        negative_prompt_ids: torch.Tensor | None,
-        negative_prompt_mask: torch.Tensor | None,


I think should recover all change like this, we should focus on feature, the fix in type should do in a separate PR if necessary

knlnguyen1802 · 2026-07-07T09:00:47Z

-        if isinstance(prompt_ids, list):
-            prompt_ids = torch.tensor(prompt_ids, device=self.device)
-        if isinstance(negative_prompt_ids, list):
-            negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)


Recover these change also

knlnguyen1802 · 2026-07-07T09:02:09Z

@Sky-Trigger Thank you for the PR, we should focus only on the feature change mention in this PR, every type correction and code linting, code refactor need to recover. We can fix them in a different PR if necessary

zhtmike · 2026-07-07T09:03:24Z

@Sky-Trigger Thank you for the PR, we should focus only on the feature change mention in this PR, every type correction and code linting, code refactor need to recover. We can fix them in a different PR if necessary

agree, pls focus on the feature only, thanks

Sky-Trigger · 2026-07-07T09:30:12Z

Thanks for the review. I may have misunderstood what you mean by “recover,” so I would like to clarify the expected implementation before making further changes.

My understanding of the review comments is that I should:

revert the type annotation changes; and
restore the following list-to-tensor conversion in _prepare_token_id_generation_context():

if isinstance(prompt_ids, list):
    prompt_ids = torch.tensor(prompt_ids, device=self.device)

if isinstance(negative_prompt_ids, list):
    negative_prompt_ids = torch.tensor(
        negative_prompt_ids,
        device=self.device,
    )

However, I am not sure how to reconcile the second change with the behavior introduced by vLLM-Omni PR #2962.

In PR #2962, the prompt embedding cache builds its key from the arguments passed to encode_prompt(). Nested lists are supported as cache-key inputs, while a call containing torch.Tensor arguments bypasses the cache.

The original verl-omni path is:

token IDs as lists
    -> convert lists to tensors in the rollout adapter
    -> call encode_prompt() with tensors
    -> prompt embedding cache bypass

The purpose of this PR is to change it to:

token IDs as lists
    -> call the cached encode_prompt() with lists
    -> cache lookup
    -> convert lists to tensors inside the underlying encode_prompt()
    -> text encoder

Therefore, if I restore the conversion in _prepare_token_id_generation_context(), the conversion will again happen before encode_prompt(), and the tokenized-prompt path will bypass the cache as before.

The type annotation changes were also made for this same reason. After delaying the tensor conversion, _prepare_token_id_generation_context() and encode_prompt() intentionally receive list inputs at runtime. I changed their annotations to describe this new input path, rather than as an unrelated typing cleanup.

Could you please clarify which implementation you would prefer?

Keep the list inputs until encode_prompt(), move the tensor conversion inside encode_prompt(), and retain the corresponding list-compatible annotations.
Keep the runtime list handling required by the cache, but revert the annotations to their previous tensor-only form and handle the annotation correction in a separate PR.
Restore both the previous annotations and the early list-to-tensor conversion. My understanding is that this would cause the prompt embedding cache to be bypassed for tokenized prompts, so please let me know if I am missing another intended cache path.

I can also revert the shared batch-size helper and any unrelated formatting or refactoring changes. I mainly need clarification on whether the delayed list-to-tensor conversion itself should remain, since it appears necessary for compatibility with PR #2962.

knlnguyen1802 · 2026-07-07T09:47:51Z

Thanks for the review. I may have misunderstood what you mean by “recover,” so I would like to clarify the expected implementation before making further changes.

My understanding of the review comments is that I should:

revert the type annotation changes; and

restore the following list-to-tensor conversion in _prepare_token_id_generation_context():
if isinstance(prompt_ids, list):
    prompt_ids = torch.tensor(prompt_ids, device=self.device)

if isinstance(negative_prompt_ids, list):
    negative_prompt_ids = torch.tensor(
        negative_prompt_ids,
        device=self.device,
    )
However, I am not sure how to reconcile the second change with the behavior introduced by vLLM-Omni PR #2962.

In PR #2962, the prompt embedding cache builds its key from the arguments passed to encode_prompt(). Nested lists are supported as cache-key inputs, while a call containing torch.Tensor arguments bypasses the cache.

The original verl-omni path is:
token IDs as lists
    -> convert lists to tensors in the rollout adapter
    -> call encode_prompt() with tensors
    -> prompt embedding cache bypass
The purpose of this PR is to change it to:
token IDs as lists
    -> call the cached encode_prompt() with lists
    -> cache lookup
    -> convert lists to tensors inside the underlying encode_prompt()
    -> text encoder
Therefore, if I restore the conversion in _prepare_token_id_generation_context(), the conversion will again happen before encode_prompt(), and the tokenized-prompt path will bypass the cache as before.

The type annotation changes were also made for this same reason. After delaying the tensor conversion, _prepare_token_id_generation_context() and encode_prompt() intentionally receive list inputs at runtime. I changed their annotations to describe this new input path, rather than as an unrelated typing cleanup.

Could you please clarify which implementation you would prefer?

Keep the list inputs until encode_prompt(), move the tensor conversion inside encode_prompt(), and retain the corresponding list-compatible annotations.

Keep the runtime list handling required by the cache, but revert the annotations to their previous tensor-only form and handle the annotation correction in a separate PR.

Restore both the previous annotations and the early list-to-tensor conversion. My understanding is that this would cause the prompt embedding cache to be bypassed for tokenized prompts, so please let me know if I am missing another intended cache path.

I can also revert the shared batch-size helper and any unrelated formatting or refactoring changes. I mainly need clarification on whether the delayed list-to-tensor conversion itself should remain, since it appears necessary for compatibility with PR #2962.

I think this PR do some unnecessary refactor and linting, let's me give a more detail review on these point

zhtmike · 2026-07-07T10:14:54Z

Not every model is in our CI, we should minimise the code change as possible to prevent any unexpected code break, thx.

Sky-Trigger · 2026-07-07T11:01:57Z

I’ll revert the unnecessary redundant changes and keep only the minimal compatibility changes required to enable prompt caching for the Qwen-Image model.

knlnguyen1802

@Sky-Trigger Please revert these thing and similar thing

knlnguyen1802 · 2026-07-07T13:35:37Z

+        prompt_ids: torch.Tensor | list[int] | list[list[int]],
+        attention_mask: torch.Tensor | list[int] | list[list[int]] | None = None,


Revert this change

knlnguyen1802 · 2026-07-07T13:38:41Z

+from verl_omni.pipelines.qwen_image_flow_grpo.common import (
+    build_img_shapes,
+    coalesce_not_none,
+    get_prompt_batch_size,
+)


All change in this file is unrelated to embedding cache, it's just convert prompt_ids to list then convert it back to tensor ?

It is true that the input is first converted to a list and later converted back to a tensor. However, the timing of these conversions is important.

After prompt_ids is converted to a list, the encode_prompt() wrapper introduced in PR #2962 can automatically apply the prompt embedding cache. The wrapper cannot cache calls when prompt_ids is a Tensor.

Therefore, the purpose of this PR is to delay the tensor conversion of prompt_ids and the other related inputs until after they have passed through the encode_prompt() wrapper, so that the prompt cache provided by PR #2962 can take effect.

knlnguyen1802 · 2026-07-07T13:39:39Z

-        if isinstance(prompt_ids, list):
-            prompt_ids = torch.tensor(prompt_ids, device=self.device)
-        if isinstance(negative_prompt_ids, list):
-            negative_prompt_ids = torch.tensor(negative_prompt_ids, device=self.device)


Why remove this ?

Do you mean that this conversion should be controlled by whether the prompt embedding cache is enabled?

For example:

if not prompt_embed_cache_enabled: if isinstance(prompt_ids, list): prompt_ids = torch.tensor(prompt_ids, device=self.device) if isinstance(negative_prompt_ids, list): negative_prompt_ids = torch.tensor( negative_prompt_ids, device=self.device, )

When the prompt embedding cache is enabled, the token IDs would remain as lists until they pass through the cached encode_prompt() call. When the cache is disabled, we would preserve the original behavior and convert them to tensors at this point.

Is this the implementation you are suggesting?

I see thanks

I will revise the implementation to preserve maximum compatibility with the existing behavior and add only the code strictly required when the prompt cache is enabled.

knlnguyen1802 · 2026-07-07T13:39:51Z

-        prompt_mask: torch.Tensor | None = None,
+        prompt_mask: torch.Tensor | list[int] | list[list[int]] | None = None,
        negative_prompt_ids: torch.Tensor | list[int] | None = None,
-        negative_prompt_mask: torch.Tensor | None = None,
+        negative_prompt_mask: torch.Tensor | list[int] | list[list[int]] | None = None,


Revert this change

knlnguyen1802 · 2026-07-07T13:40:16Z

-        prompt_ids: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
+        prompt_ids: torch.Tensor | list[int] | list[list[int]],
+        attention_mask: torch.Tensor | list[int] | list[list[int]] | None = None,


Revert this

knlnguyen1802 · 2026-07-07T13:42:27Z

+            if isinstance(prompt_ids, list):
+                prompt_ids = torch.tensor(prompt_ids, device=self.device)
+            if isinstance(attention_mask, list):
+                attention_mask = torch.tensor(attention_mask, device=self.device)


If you want to refactor this as common, then can do it in a separate PR

knlnguyen1802 · 2026-07-07T13:42:42Z

-        prompt_mask: torch.Tensor | None = None,
+        prompt_mask: torch.Tensor | list[int] | list[list[int]] | None = None,
        negative_prompt_ids: torch.Tensor | list[int] | None = None,
-        negative_prompt_mask: torch.Tensor | None = None,
+        negative_prompt_mask: torch.Tensor | list[int] | list[list[int]] | None = None,


Revert this

knlnguyen1802 · 2026-07-07T13:42:56Z

-        prompt_ids: torch.Tensor,
-        attention_mask: torch.Tensor | None = None,
+        prompt_ids: torch.Tensor | list[int] | list[list[int]],
+        attention_mask: torch.Tensor | list[int] | list[list[int]] | None = None,


Revert this

knlnguyen1802 · 2026-07-07T13:43:07Z

-        prompt_mask: torch.Tensor | None = None,
+        prompt_mask: torch.Tensor | list[int] | list[list[int]] | None = None,
        negative_prompt_ids: torch.Tensor | list[int] | None = None,
-        negative_prompt_mask: torch.Tensor | None = None,
+        negative_prompt_mask: torch.Tensor | list[int] | list[list[int]] | None = None,


Same, revert this

knlnguyen1802 · 2026-07-07T13:43:43Z

+def get_prompt_batch_size(prompt_ids: torch.Tensor | list[int] | list[list[int]]) -> int:
+    if isinstance(prompt_ids, torch.Tensor):
+        return prompt_ids.shape[0] if prompt_ids.ndim == 2 else 1
+    return len(prompt_ids) if prompt_ids and isinstance(prompt_ids[0], list) else 1


I think it's a good idea to refactor this as common feature, but unrelated to this PR

Please do it in a separate PR, and revert all related things

Sky-Trigger · 2026-07-07T16:28:47Z

I have completed the changes. Please review them again. @knlnguyen1802 @zhtmike
This time, I only added the necessary support for enabling the prompt cache while preserving compatibility with the previous implementation. Note that this change does not modify the parameter types passed within the function. Therefore, the type annotation may allow List[int], but the actual input remains a tensor.

Sky-Trigger · 2026-07-08T03:39:13Z

I still believe that the changes introduced in my first version are necessary, specifically the changes to the accepted input types and the decision to defer the tensor conversion of prompt_ids and the corresponding masks so that they can work correctly with the prompt-cache implementation in vLLM-Omni.

I think there may have been some misunderstanding regarding my previous explanation. The changes described above are required for prompt caching to function correctly. Simply adding configuration options such as enable_prompt_embed_cache is not sufficient, because the current early tensor conversion prevents the inputs from entering the prompt-cache path correctly.

In addition, these changes are mainly made in the adapter layer and common.py, so they are generally applicable to all algorithms rather than being specific to the Qwen-Image model. The implementation is intended to preserve the existing behavior when prompt caching is disabled, so it should not break the current functionality.

knlnguyen1802 · 2026-07-08T03:42:45Z

I still believe that the changes introduced in my first version are necessary, specifically the changes to the accepted input types and the decision to defer the tensor conversion of prompt_ids and the corresponding masks so that they can work correctly with the prompt-cache implementation in vLLM-Omni.

In addition, these changes are mainly made in the adapter layer and common.py, so they are generally applicable to all algorithms rather than being specific to the Qwen-Image model. The implementation is intended to preserve the existing behavior when prompt caching is disabled, so it should not break the current functionality.

You can do 2 separate PR for them, refactor PR usually easy to review. Because our CI does not cover too broad yet so decouple PR is prefered

Sky-Trigger · 2026-07-08T03:47:39Z

Please take another look at the latest changes. This version now contains only the minimum modifications necessary to enable the required support only for Qwen-Image.

Signed-off-by: Trigger <129651635+Sky-Trigger@users.noreply.github.com>

knlnguyen1802 · 2026-07-08T03:55:33Z

Please take another look at the latest changes. This version now contains only the minimum modifications necessary to enable the required support only for Qwen-Image.

I think we can do it for Qwen-Image first, to see any regression, later can expand for another model

knlnguyen1802 · 2026-07-08T03:55:50Z

Overall LGTM, cc @zhtmike for final review

zhtmike · 2026-07-08T04:00:28Z

@wtomin and @AndyZhou952 , for the DPO / DiffusionNFT change, PTAL

wtomin · 2026-07-08T06:41:12Z

        custom_prompt: OmniCustomPrompt = {"prompt_token_ids": prompt_ids}
        if prompt_mask is not None:
-            custom_prompt["prompt_mask"] = prompt_mask
+            custom_prompt["prompt_mask"] = (


If unconditionally converting prompt_mask from tensor to list, will it cause some errors when calling attention_mask.ndim? Please double check it. A better solution might be converting prompt_mask to list only wen prompt embedding is enabled.

zhtmike · 2026-07-08T06:46:40Z

After an offline discussion with @knlnguyen1802, we agreed that it may create too much development burden if we are forced to convert tensors to lists just to enable the prompt cache. It may be better to fix this on the vLLM‑Omni side instead. Let us hold for this PR now

Sky-Trigger · 2026-07-08T06:50:55Z

ok tks，pls let me know if any further changes are needed.

Sky-Trigger requested review from SamitHuang and zhtmike as code owners June 29, 2026 06:21

gemini-code-assist Bot reviewed Jun 29, 2026

View reviewed changes

zhtmike requested a review from knlnguyen1802 June 29, 2026 06:35

knlnguyen1802 reviewed Jul 7, 2026

View reviewed changes

Sky-Trigger requested a review from knlnguyen1802 July 7, 2026 09:33

knlnguyen1802 reviewed Jul 7, 2026

View reviewed changes

[rollout] limit prompt cache support to Qwen-Image

f7cf431

Sky-Trigger force-pushed the prompt-embed-cache-adapters branch from 67154d5 to f7cf431 Compare July 7, 2026 16:23

Sky-Trigger changed the title ~~[rollout, diffusion] fix: support prompt embedding cache for tokenized prompts~~ [rollout, diffusion] fix: support prompt embedding cache for Qwen-Image Jul 7, 2026

Sky-Trigger requested a review from knlnguyen1802 July 8, 2026 01:04

Merge branch 'main' into prompt-embed-cache-adapters

9f01c44

Update common.py

10daba2

Signed-off-by: Trigger <129651635+Sky-Trigger@users.noreply.github.com>

zhtmike requested review from AndyZhou952 and wtomin July 8, 2026 03:59

wtomin reviewed Jul 8, 2026

View reviewed changes

		prompt_ids: torch.Tensor \| list[int] \| list[list[int]],
		attention_mask: torch.Tensor \| list[int] \| list[list[int]] \| None = None,

Uh oh!

Conversation

Sky-Trigger commented Jun 29, 2026

What does this PR do?

Checklist Before Starting

Test

API and Usage Example

Design & Code Changes

Prompt cache boundary

Adapter changes

Configuration changes

Checklist Before Submitting

AI Assistance

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

zhtmike commented Jun 29, 2026

Uh oh!

knlnguyen1802 commented Jul 3, 2026

Uh oh!

Sky-Trigger commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

knlnguyen1802 commented Jul 6, 2026

Uh oh!

Sky-Trigger commented Jul 7, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

knlnguyen1802 commented Jul 7, 2026

Uh oh!

zhtmike commented Jul 7, 2026

Uh oh!

Sky-Trigger commented Jul 7, 2026

Uh oh!

knlnguyen1802 commented Jul 7, 2026

Uh oh!

zhtmike commented Jul 7, 2026

Uh oh!

Sky-Trigger commented Jul 7, 2026

Uh oh!

knlnguyen1802 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Sky-Trigger commented Jul 7, 2026

Uh oh!

Sky-Trigger commented Jul 8, 2026

Uh oh!

Sky-Trigger commented Jul 6, 2026 •

edited

Loading